package com.leetcode.题库;

import java.util.Arrays;

/**
 * @author 覃国强
 * @date 2022-03-16 11:27
 */
public class 题库_0498_对角线遍历 {

  public static void main(String[] args) {
    Solution solution = new Solution();
    int[][] mat = {{1,2,3},{4,5,6},{7,8,9}};
    int[] diagonalOrder = solution.findDiagonalOrder(mat);
    System.out.println(Arrays.toString(diagonalOrder));
  }

  static
  class Solution {

    private static final String TREND_UP = "UP";
    private static final String TREND_DOWN = "DOWN";

    public int[] findDiagonalOrder(int[][] mat) {
      if (mat == null || mat.length == 0 || mat[0] == null || mat[0].length == 0) {
        return new int[0];
      }
      int m = mat.length;
      int n = mat[0].length;
      int[] answer = new int[m * n];

      int index = 0;
      int i = 0;
      int j = 0;
      String trend = TREND_UP;

      do {
        answer[index++] = mat[i][j];
        if (i == m - 1 && j == n - 1) {
          return answer;
        } else if (TREND_UP.equals(trend)) {
          if (i == 0 || j == n - 1) {
            int[] next = getNext(i, j, m, n, trend);
            i = next[0];
            j = next[1];
            trend = TREND_DOWN;
          } else {
            i -= 1;
            j += 1;
          }
        } else {
          if (i == m - 1 || j == 0) {
            int[] next = getNext(i, j, m, n, trend);
            i = next[0];
            j = next[1];
            trend = TREND_UP;
          } else {
            i += 1;
            j -= 1;
          }
        }
      } while (true);
    }

    private int[] getNext(int i, int j, int m, int n, String trend) {
      if (TREND_UP.equals(trend)) {
        if (j == n - 1) {
          // i = i + 1, j = n - 1
          return new int[]{i + 1, n - 1};
        }
        if (i == 0) {
          // i = 0, j = j + 1
          return new int[]{0, j + 1};
        }
      }
      if (TREND_DOWN.equals(trend)) {
        if (i == m - 1) {
          // i = m - 1, j = j + 1
          return new int[]{m - 1, j + 1};
        }
        if (j == 0) {
          // i = i + 1, j = 0
          return new int[]{i + 1, 0};
        }
      }
      return new int[2];
    }

  }

}
